# 使用 RNN 处理语句

import torch

input_size = 4
hidden_size = 4
# batch_size 只用于生成初始的隐层输入 h0 有关
batch_size = 1
num_layers = 1
seq_len = 5

# 词汇表
idx2char = ['e', 'h', 'l', 'o']

# 训练集和样本
x_data = [1, 0, 2, 2, 3]
y_data = [3, 1, 2, 3, 2]

# 独热向量矩阵，用于查找相应的独热向量
one_hot_lookup = [[1, 0, 0, 0],
                  [0, 1, 0, 0],
                  [0, 0, 1, 0],
                  [0, 0, 0, 1]]

# 将输入样本的特征使用独热向量表示
x_one_hot = [one_hot_lookup[x] for x in x_data]
# 修改输入样本和标签的特征维度
# inputs = torch.Tensor(x_one_hot).view(-1, batch_size, input_size)
inputs = torch.Tensor(x_one_hot).view(seq_len, batch_size, input_size)
# 这里需要labels是要给一维，因为是一次性传给RNN，不再需要抽出单个元素比较
labels = torch.LongTensor(y_data)


# labels = torch.LongTensor(y_data)


# 设计模型
class Model(torch.nn.Module):
    def __init__(self, input_size, hidden_size, batch_size, num_layers):
        super(Model, self).__init__()
        self.num_layers = num_layers
        self.batch_size = batch_size
        self.input_size = input_size
        self.hidden_size = hidden_size
        # 定义 RNN Cell 结构
        self.rnn = torch.nn.RNN(input_size=self.input_size,
                                hidden_size=self.hidden_size,
                                num_layers=self.num_layers)

    # 前馈计算，。传入参数 inputs为输入，hidden为隐层输入
    def forward(self, input):
        hidden = torch.zeros(self.num_layers,
                             self.batch_size,
                             self.hidden_size)
        out, _ = self.rnn(input, hidden)
        return out.view(-1, self.hidden_size)


# 模型实例
model = Model(input_size, hidden_size, batch_size, num_layers)

# 定义损失和优化器
criterion = torch.nn.CrossEntropyLoss()
# 这里改用一个新的算法 Adam 作为优化算法，学习率改为0.1加速收敛速度
optimizer = torch.optim.Adam(model.parameters(), lr=0.1)

# 训练模型
for epoch in range(15):
    # 梯度清0
    optimizer.zero_grad()
    # 获得输出
    outputs = model(inputs)
    # 计算损失
    loss = criterion(outputs, labels)
    # 反向传播
    loss.backward()
    # 参数更新
    optimizer.step()
    _, idx = outputs.max(dim=1)
    idx = idx.data.numpy()
    print('Predicted: ', ''.join([idx2char[x] for x in idx]), end='')
    print(', Epoch [%d/15] loss = %.3f' % (epoch + 1, loss.item()))
